﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading.Tasks;
using SharpRBM.Core.Enumerations;

namespace SharpRBM.Core.LearningModules
{
    public class ArrayBasedLearningModule : ILearningModule
    {
        private const bool MultiThreaded = true;
        private float[] _rawTestCases;
        private float[] _detectorError;
        private int _testCount;
        private InterLayerWeights _weights;
        private List<object> _locks;
        public void Prepare(InterLayerWeights weights, List<Layer> inputs)
        {
            _weights = weights;
            _rawTestCases = LearningModuleHelper.CreateRawTestCases(inputs, _weights.LowerLayerSize);
            _testCount = inputs.Count;
            _detectorError = new float[weights.UpperLayerSize * weights.LowerLayerSize];
            _locks = new List<object>();
            Helper.FillList(_locks, weights.LowerLayerSize, ()=>new object());
        }

        public void ActivateLowerToUpper(Layer lower, Layer upper, ActivationMethod method)
        {
            _weights.ActivateLowerToUpper(lower, upper, method); 
        }

        public void ActivateUpperToLower(Layer lower, Layer upper)
        {
            _weights.ActivateUpperToLower(lower, upper);
        }
        
        public TrainingError Train()
        {
            ClearDetectorErrors(_weights.LowerLayerSize, _weights.UpperLayerSize);

            float reconstructionError = 0;

            ParallelFor(
                MultiThreaded,
                0,
                _testCount,
                testCase =>
                {
                    float errorPart =
                        TrainOnSingleCase(
                        _rawTestCases,
                        _weights.Weights,
                        _detectorError,
                        testCase,
                        _weights.LowerLayerSize,
                        _weights.UpperLayerSize,
                        _testCount);

                    lock (_locks[testCase % _weights.LowerLayerSize])
                    {
                        reconstructionError += errorPart;
                    }
                });

            float epsilon = _weights.GetAdjustedAndScaledTrainingRate(_testCount);
            UpdateWeights(_weights.Weights, _weights.LowerLayerSize, _weights.UpperLayerSize, _detectorError, epsilon);
            TrainingError trainingError = new TrainingError(_detectorError.Sum(val => Math.Abs(val)), reconstructionError);
            _weights.RegisterLastTrainingError(trainingError);
            return trainingError;
        }

        public static void ParallelFor(bool parallel, int from, int to, Action<int> action)
        {
            if (parallel)
            {
                Parallel.For(from, to, action);
            }
            else
            {
                for (int i = from; i < to; i++)
                {
                    action(i);
                }
            }
        }

        public void Dispose()
        {
            // Not strictly required...
            _rawTestCases = null;
        }

        private float TrainOnSingleCase(
            float[] rawTestCases,
            float[] weights,
            float[] detectorErrors,
            int testCase,
            int lowerCount,
            int upperCount,
            int testCaseCount)
        {
            float[] model = new float[upperCount];
            float[] reconstructed = new float[lowerCount];
            float[] reconstructedModel = new float[upperCount];
            int rawTestCaseOffset = testCase * lowerCount;

            // Model
            ActivateLowerToUpperBinary(rawTestCases, lowerCount, rawTestCaseOffset, model, upperCount, weights);

            // Reconstruction
            ActivateUpperToLower(reconstructed, lowerCount, model, upperCount, weights);

            // Reconstruction model
            ActivateLowerToUpper(reconstructed, lowerCount, 0, reconstructedModel, upperCount, weights);

            // Accumulate detector errors
            return AccumulateErrors(rawTestCases, lowerCount, rawTestCaseOffset, model, upperCount, reconstructed, reconstructedModel, detectorErrors);
        }

        private void ClearDetectorErrors(int lowerCount, int upperCount)
        {
            for (int lower = 0; lower < lowerCount; lower++)
            {
                int offset = upperCount * lower;
                for (int upper = 0; upper < upperCount; upper++)
                {
                    _detectorError[offset + upper] = 0;
                }
            }
        }

        private void ActivateLowerToUpper(float[] lowerValues, int lowerCount, int lowerOffset, float[] upperValues, int upperCount, float[] weights)
        {
            for (int upper = 0; upper < upperCount; upper++)
            {
                float val = 0;
                for (int lower = 0; lower < lowerCount; lower++)
                {
                    val += weights[upperCount * lower + upper] * lowerValues[lower + lowerOffset];
                }

                upperValues[upper] = 1.0F / (1.0F + (float)Math.Exp(-val));
            }
            upperValues[0] = 1F;
        }

        private void ActivateLowerToUpperBinary(float[] lowerValues, int lowerCount, int lowerOffset, float[] upperValues, int upperCount, float[] weights)
        {
            for (int upper = 0; upper < upperCount; upper++)
            {
                float val = 0;
                for (int lower = 0; lower < lowerCount; lower++)
                {
                    float activation = lowerValues[lower + lowerOffset];
                    double random;
                    lock (RBMCommon.Random)
                    {
                        random = RBMCommon.Random.NextDouble();
                    }
                    activation = random <= activation ? activation : 0;
                    val += weights[upperCount * lower + upper] * activation;
                }

                upperValues[upper] = 1.0F / (1.0F + (float)Math.Exp(-val));
            }
            upperValues[0] = 1F;
        }

        private void ActivateUpperToLower(float[] lowerValues, int lowerCount, float[] upperValues, int upperCount, float[] weights)
        {
            for (int lower = 0; lower < lowerCount; lower++)
            {
                float val = 0;
                int offset = upperCount * lower;
                for (int upper = 0; upper < upperCount; upper++)
                {
                    val += weights[offset + upper] * upperValues[upper];
                }

                lowerValues[lower] = 1.0F / (1.0F + (float)Math.Exp(-val));
            }
            lowerValues[0] = 1F;
        }

        private float AccumulateErrors(
            float[] rawTestCases,
            int lowerCount,
            int rawTestCaseOffset,
            float[] model,
            int upperCount,
            float[] reconstructed,
            float[] reconstructedModel,
            float[] detectorErrors)
        {
            float reconstructedError = 0;
            float[] errorRow = new float[upperCount];
            for (int lower = 0; lower < lowerCount; lower++)
            {
                int errorOffset = upperCount * lower;
                for (int upper = 0; upper < upperCount; upper++)
                {
                    errorRow[upper] =
                        // What the model should believe in
                        rawTestCases[rawTestCaseOffset + lower] * model[upper] +
                        // What the model actually believes in
                        -reconstructed[lower] * reconstructedModel[upper];
                }

                // Minimize the amount of locking by locking one row update at a time instead of locking
                // for each write
                lock (_locks[lower])
                {
                    for (int upper = 0; upper < upperCount; upper++)
                    {
                        detectorErrors[errorOffset + upper] -= errorRow[upper];
                    }
                }

                reconstructedError += Math.Abs(rawTestCases[rawTestCaseOffset + lower] - reconstructed[lower]);
            }

            return reconstructedError;
        }

        private void UpdateWeights(float[] weights, int lowerCount, int upperCount, float[] detectorErrors, float epsilon)
        {
            for (int lower = 0; lower < lowerCount; lower++)
            {
                int offset = upperCount * lower;
                for (int upper = 0; upper < upperCount; upper++)
                {
                    int pos = offset + upper;
                    weights[pos] -= detectorErrors[pos] * epsilon;
                }
            }
        }
    }
}
